from src.document_parser import DocumentParser
from src.image_handler import ImageHandler
from src.content_processor import ContentProcessor
import yaml
import json
import os
from pathlib import Path

def load_config():
    """加载配置文件"""
    with open('../config/config.yaml', 'r', encoding='utf-8') as f:
        return yaml.safe_load(f)

def process_document(file_path: str, source_name: str):
    """处理单个文档
    
    Args:
        file_path: 文档路径
        source_name: 文档英文名称
    """
    config = load_config()
    
    # 1. 初始化各个组件
    doc_parser = DocumentParser(config)
    image_handler = ImageHandler(config)
    content_processor = ContentProcessor(config)
    
    print(f"\n开始处理文档: {file_path}")
    
    try:
        # 创建输出目录
        output_dir = Path('output')
        output_dir.mkdir(exist_ok=True)
        
        # 生成输出文件名（不带扩展名）
        base_name = Path(file_path).stem
        output_base = output_dir / f"{source_name}_{base_name}"
        
        # 检查是否存在已处理的markdown文件
        markdown_file = output_base.with_suffix('.md')
        if markdown_file.exists():
            print(f"发现已处理的Markdown文件: {markdown_file}")
            with open(markdown_file, 'r', encoding='utf-8') as f:
                markdown_content = f.read()
            print(f"   - 已读取现有Markdown文件，包含 {len(markdown_content)} 字符")
        else:
            # 1. 解析文档并生成markdown格式txt
            print("1. 解析文档...")
            markdown_content = doc_parser.parse_to_markdown(file_path, image_handler, source_name)
            print(f"   - 已生成Markdown文件，包含 {len(markdown_content)} 字符")
            
            # 保存markdown文件
            with open(markdown_file, 'w', encoding='utf-8') as f:
                f.write(markdown_content)
            print(f"   - Markdown文件已保存到: {markdown_file}")
        
        # 2. 语义分段并生成QA对
        print("\n2. 生成问答对...")
        qa_pairs = content_processor.generate_qa_pairs(markdown_content, source_name)
        print(f"   - 生成问答对: {len(qa_pairs)} 个")
        
        # 3. 保存QA结果
        qa_file = output_base.with_suffix('.json')
        with open(qa_file, 'w', encoding='utf-8') as f:
            json.dump({
                'source_name': source_name,
                'original_file': str(file_path),
                'markdown_file': str(markdown_file),
                'qa_pairs': qa_pairs
            }, f, ensure_ascii=False, indent=2)
        
        print(f"\n处理完成!")
        print(f"- Markdown文件: {markdown_file}")
        print(f"- QA结果文件: {qa_file}")
        
    except Exception as e:
        print(f"\n处理出错: {str(e)}")
        raise e

if __name__ == "__main__":
    # 测试文件路径
    file_path = r"E:\Desktop\操作手册知识库\攸信uMOM  6.0 配置中心-品质配置-不良管理操作手册_V1.0-20241219.pdf"
    source_name = "testDoc"
    
    process_document(file_path, source_name) 